from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session

from app import crud, models, schemas
from app.api import deps
from app.models.credit_transaction import TransactionType

router = APIRouter()

# A mock cost for creating a job
JOB_COST = 10

@router.post("/", response_model=schemas.GenerationJob)
def create_job(
    *,
    db: Session = Depends(deps.get_db),
    job_in: schemas.GenerationJobCreate,
    current_user: models.User = Depends(deps.get_current_user),
) -> any:
    """
    Create new generation job.
    This endpoint will deduct credits from the user's balance.
    """
    if current_user.credits_balance < JOB_COST:
        raise HTTPException(
            status_code=status.HTTP_402_PAYMENT_REQUIRED,
            detail="Insufficient credits. Please recharge.",
        )

    # Deduct credits and create the job in a single transaction
    try:
        crud.update_user_credits(
            db, 
            user_id=current_user.id, 
            amount=-JOB_COST, 
            trans_type=TransactionType.GENERATION_FEE,
            description=f"AI生成任务费用: {job_in.prompt[:50]}..."
        )
        job = crud.create_generation_job(db=db, job_in=job_in, user_id=current_user.id)
    except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))

    # 触发Celery异步任务
    from app.tasks import generate_image_task, generate_video_task
    
    # 根据任务类型选择不同的处理任务
    task_type = job_in.params.get("type", "image") if job_in.params else "image"
    
    if task_type == "video":
        task = generate_video_task.delay(job_id=job.id)
    else:
        task = generate_image_task.delay(job_id=job.id)
        
    # 可以记录任务ID到数据库中，用于查询任务状态
    # update_job(db, db_obj=job, obj_in={"celery_task_id": task.id})

    return job


@router.get("/{job_id}", response_model=schemas.GenerationJob)
def read_job(
    *,
    db: Session = Depends(deps.get_db),
    job_id: int,
    current_user: models.User = Depends(deps.get_current_user),
) -> any:
    """
    Get generation job by ID.
    """
    job = crud.get_job(db=db, job_id=job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    if job.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="Not enough permissions")
    return job


@router.post("/{job_id}/publish", response_model=schemas.GenerationJob)
def publish_job(
    *,
    db: Session = Depends(deps.get_db),
    job_id: int,
    publish_in: schemas.GenerationJobUpdate,
    current_user: models.User = Depends(deps.get_current_user),
) -> any:
    """
    Publish a job to the gallery.
    """
    job = crud.get_job(db=db, job_id=job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    if job.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="Not enough permissions")
    updated_job = crud.update_job(db=db, db_obj=job, obj_in=publish_in)
    return updated_job


@router.get("/", response_model=list[schemas.GenerationJob])
def list_user_jobs(
    db: Session = Depends(deps.get_db),
    current_user: models.User = Depends(deps.get_current_user),
    skip: int = 0,
    limit: int = 100,
) -> any:
    """
    Get current user's jobs.
    """
    jobs = crud.get_user_jobs(db=db, user_id=current_user.id, skip=skip, limit=limit)
    return jobs
